﻿Imports Encryptor.My
Imports Encryptor.My.Templates
Imports System.Text

Public Class frmCrypto
    Private Const EncFileExt As String = ".aesx"
    Private myDocs$ = My.Computer.FileSystem.SpecialDirectories.MyDocuments
    Dim WFEcipherText As Byte()
    Dim deciphertext As Byte()
    Dim failed As Int32 = 0
    Dim desktop As String = My.Computer.FileSystem.SpecialDirectories.Desktop

    Private Sub frmCrypto_Load(ByVal sender As System.Object, ByVal e As System.EventArgs) Handles MyBase.Load
        If Application.IsNetworkDeployed Then
            Me.Text = "AES Encryptor " & My.Application.Deployment.CurrentVersion.ToString()
        Else
            Me.Text = "AES Encryptor " & My.Application.Info.Version.ToString()

        End If

        If My.Application.CommandLineArgs.Count = 4 Then
            Me.Hide()
            Dim type = My.Application.CommandLineArgs(0)
            Dim file = My.Application.CommandLineArgs(1)
            Dim pSwitch = My.Application.CommandLineArgs(2)
            Dim password = My.Application.CommandLineArgs(3)

            If pSwitch = "-p" Then
                If password = "default" Then

                    If type = "-e" Then
                        Call EncryptFile(file, MySettings.Default.password)
                    ElseIf type = "-d" Then
                        Call DecryptFile(file, MySettings.Default.password)
                    End If

                Else

                    If type = "-e" Then
                        Call EncryptFile(file, password)
                    ElseIf type = "-d" Then
                        Call DecryptFile(file, password)
                    End If

                End If
            End If

        End If

        If MySettings.Default.password = "" Or IsNothing(MySettings.Default.password) Then
            MsgBox("The password is currently set to 'newuser'.")
        End If

        If MySettings.Default.logOption = "" Or IsNothing(MySettings.Default.logOption) Then
            MySettings.Default.logOption = "off"
        End If
        ofdEncWholeFile.InitialDirectory = myDocs
    End Sub

    Private Sub ofdEncWholeFile_FileOk(ByVal sender As System.Object, ByVal e As System.ComponentModel.CancelEventArgs) Handles ofdEncWholeFile.FileOk
        Dim tmp(2) As String
        Dim lvi As ListViewItem
        For i As Int32 = 0 To (ofdEncWholeFile.FileNames.Length - 1) 'Add the ENCRYPT property to all the selected files
            tmp(0) = ofdEncWholeFile.FileNames(i.ToString)
            tmp(1) = "ENCRYPT"
            lvi = New ListViewItem(tmp)
            lvwLoad.Items.Add(lvi)
        Next
    End Sub

    Private Sub btnEnc_Click(ByVal sender As System.Object, ByVal e As System.EventArgs) Handles btnEnc.Click
        dlgCipherText.ShowDialog()
    End Sub

    Private Sub ofdDecFile_Disposed(ByVal sender As Object, ByVal e As System.EventArgs) Handles ofdDecFile.Disposed
        Me.Focus()
    End Sub

    Private Sub ofdDecFile_FileOk(ByVal sender As System.Object, ByVal e As System.ComponentModel.CancelEventArgs) Handles ofdDecFile.FileOk
        Dim tmp(2) As String
        Dim lvi As ListViewItem
        For i As Int32 = 0 To (ofdDecFile.FileNames.Length - 1) 'Add the DECRYPT property to all the selected files
            tmp(0) = ofdDecFile.FileNames(i.ToString)
            tmp(1) = "DECRYPT"
            lvi = New ListViewItem(tmp)
            lvwLoad.Items.Add(lvi)
        Next
    End Sub

    Private Sub btnChPass_Click(ByVal sender As System.Object, ByVal e As System.EventArgs) Handles btnChPass.Click
        dlgChPass.ShowDialog()
    End Sub

    Private Sub EncryptFile(ByVal filepath As String, passkey As String)
        Dim safeName As String = filepath.Substring(filepath.LastIndexOf("\") + 1)
        Dim plainTextBytes = Computer.FileSystem.ReadAllBytes(filepath)
        Dim encAes As New SimpleAes(passkey)
        Dim WFEcipherText As Byte() = encAes.EncryptData(plainTextBytes)
        Computer.FileSystem.WriteAllBytes(filepath, WFEcipherText, False)
        Computer.FileSystem.RenameFile(filepath, safeName & EncFileExt)
        Console.WriteLine("File Encrypted Successfully")
    End Sub

    Private Sub DecryptFile(ByVal filepath As String, passkey As String)
        Dim safeName As String = filepath.Substring(filepath.LastIndexOf("\") + 1)
        Dim cipherTextBytes = Computer.FileSystem.ReadAllBytes(filepath)
        Dim decAes As New SimpleAes(passkey)
        Dim plainTextBytes As Byte() = decAes.DecryptData(cipherTextBytes)
        Computer.FileSystem.WriteAllBytes(filepath, plainTextBytes, False)
        Dim originalName$ = filepath.Remove(filepath.LastIndexOf("."c))
        originalName = originalName.Substring(originalName.LastIndexOf("\"c) + 1)
        Computer.FileSystem.RenameFile(filepath, originalName)
        Console.WriteLine("File Decrypted Successfully")
    End Sub

    Private Sub btnLoadDec_Click(sender As System.Object, e As System.EventArgs) Handles btnLoadDec.Click
        ofdDecFile.ShowDialog()
    End Sub

    Private Sub btnLoadEnc_Click(sender As System.Object, e As System.EventArgs) Handles btnLoadEnc.Click
        ofdEncWholeFile.ShowDialog()
    End Sub

    Private Sub btnProc_Click(sender As System.Object, e As System.EventArgs) Handles btnProc.Click
        If lvwLoad.Items.Count < 1 Then 'Check for an empty list 
            MsgBox("There is nothing to process.", MsgBoxStyle.OkOnly + MsgBoxStyle.Exclamation)
            Exit Sub
        End If

        pbarMain.Maximum = lvwLoad.Items.Count

        If LookForDecrypt() = True Then 'Check for any DECRYPT operations loaded in the list
            Select Case logDecrypt.ShowDialog()
                Case Windows.Forms.DialogResult.OK 'If the user clicked OK and the password is good, exit the select block and run the decryption sub
                    Exit Select
                Case Windows.Forms.DialogResult.Abort 'If the user clicked OK and the password is bad, exit the whole process with an error message
                    MsgBox("Incorrect password.", MsgBoxStyle.OkOnly + MsgBoxStyle.Exclamation)
                    Exit Sub
                Case Windows.Forms.DialogResult.Cancel 'If the user clicked Cancel, exit the whole process with an info message
                    MsgBox("Process canceled.", MsgBoxStyle.OkOnly + MsgBoxStyle.Information)
                    Exit Sub
            End Select

            Call ProcDecrypt() 'Call the decryption sub

        End If

        If LookForEncrypt() = True Then
            Call ProcEncrypt()
        End If


        Dim total = lvwLoad.Items.Count - failed

        If total >= 1 AndAlso My.MySettings.Default.logOption = "on" Then
            GenLog()
        ElseIf total >= 1 AndAlso My.MySettings.Default.logOption = "prompt" Then
            If MsgBox("Create a log for this process?", MsgBoxStyle.Question + MsgBoxStyle.YesNo) = MsgBoxResult.Yes Then
                GenLog()
            End If
        End If

        If total > 1 Then
            MsgBox(total & " items successfully processed." & vbCrLf & failed & " error(s)", MsgBoxStyle.OkOnly + MsgBoxStyle.Information)
        ElseIf total = 1 Then
            MsgBox(total & " item successfully processed." & vbCrLf & failed & " error(s)", MsgBoxStyle.OkOnly + MsgBoxStyle.Information)
        ElseIf total < 1 Then
            MsgBox("None of the items could be processed." & vbCrLf & failed & " error(s)", MsgBoxStyle.Exclamation + MsgBoxStyle.OkOnly)
        End If

        pbarMain.Value = 0
        failed = 0
        lvwLoad.Items.Clear()
    End Sub

    Private Sub btnRemove_Click(sender As System.Object, e As System.EventArgs) Handles btnRemove.Click
        If lvwLoad.SelectedItems.Count < 1 Then
            MsgBox("You have not selected anything to remove.", MsgBoxStyle.OkOnly + MsgBoxStyle.Exclamation)
        Else
            For Each lvi As ListViewItem In lvwLoad.SelectedItems
                lvwLoad.Items.Remove(lvi)
            Next
        End If
    End Sub

    Private Function LookForDecrypt() As Boolean
        Dim procList(lvwLoad.Items.Count) As String
        For i As Int32 = 0 To (lvwLoad.Items.Count - 1) 'Enumerate all the process types
            procList(i) = lvwLoad.Items(i).SubItems(1).Text
        Next

        'Check the list for DECRYPT operations
        If procList.Contains("DECRYPT") Then 'If there are, bring up the password box and get the DECRYPTs out of the way
            Return True
        Else 'If not, just run the ENCRYPTs
            Return False
        End If
    End Function

    Private Sub ProcEncrypt()
        'Fire up the variables
        Dim decAes As New SimpleAes(MySettings.Default.password)
        Dim plainTextBytes As Byte() = {}
        Dim enclist(GetTypeCount(0) - 1) As String
        Dim count% = 0

        'Get a list of all the ENCRYPTs
        For i As Int32 = 0 To (lvwLoad.Items.Count - 1) 'Iterate through all the process types, weeding out all but the ENCRYPTs
            If lvwLoad.Items(i).SubItems(1).Text = "ENCRYPT" Then
                enclist(count) = (lvwLoad.Items.Item(i).Text)
                count += 1
            End If
        Next


        'Encrypt the files
        Select Case MySettings.Default.obfuscate
            Case True
                For j = 0 To (enclist.Length - 1)
                    plainTextBytes = Computer.FileSystem.ReadAllBytes(enclist(j).ToString())
                    WFEcipherText = decAes.EncryptData(plainTextBytes)
                    Computer.FileSystem.WriteAllBytes(enclist(j).ToString(), WFEcipherText, False)
                    Computer.FileSystem.RenameFile(enclist(j).ToString(), ObfucateFilename(enclist(j).ToString.Substring(enclist(j).LastIndexOf("\") + 1)) & EncFileExt)
                    pbarMain.Increment(1)
                Next
            Case False
                For j = 0 To (enclist.Length - 1)
                    plainTextBytes = Computer.FileSystem.ReadAllBytes(enclist(j).ToString())
                    WFEcipherText = decAes.EncryptData(plainTextBytes)
                    Computer.FileSystem.WriteAllBytes(enclist(j).ToString(), WFEcipherText, False)
                    Computer.FileSystem.RenameFile(enclist(j).ToString(), enclist(j).ToString.Substring(enclist(j).LastIndexOf("\") + 1) & EncFileExt)
                    pbarMain.Increment(1)
                Next
        End Select


        Return
    End Sub

    Private Sub ProcDecrypt()
        'Fire up the variables
        Dim encAes As New SimpleAes(MySettings.Default.password)
        Dim originalName As String = ""
        Dim cipherTextBytes As Byte() = {}
        Dim declist(GetTypeCount(1) - 1) As String
        Dim count% = 0

        'Get a list of all the DECRYPTs
        For i As Int32 = 0 To (lvwLoad.Items.Count - 1) 'Iterate through all the process types, weeding out all but the DECRYPTs
            If lvwLoad.Items(i).SubItems(1).Text = "DECRYPT" Then
                declist(count) = (lvwLoad.Items.Item(i).Text)
                count += 1
            End If
        Next

        Dim errorReport As New ErrLogFile 'Prepare for errors
        Dim errorText As String = errorReport.TransformText()

        'Decrypt the files
        For j = 0 To (declist.Length - 1)
            cipherTextBytes = Computer.FileSystem.ReadAllBytes(declist(j).ToString())
            Try
                deciphertext = encAes.DecryptData(cipherTextBytes)
                Computer.FileSystem.WriteAllBytes(declist(j).ToString(), deciphertext, False)

                originalName = declist(j).ToString().Remove(declist(j).ToString().LastIndexOf("."c)) 'Get the original name back
                originalName = originalName.Substring(originalName.LastIndexOf("\"c) + 1)

                If IsBase64(originalName) Then
                    originalName = DeobfuscateFilename(originalName)
                End If

                Computer.FileSystem.RenameFile(declist(j).ToString(), originalName)
            Catch ex As CryptographicException
                MsgBox("The password has been changed since encryption of this file. The data in file " & vbCrLf & "'" & declist(j).ToString() & "'" & vbCrLf & " cannot be decrypted. Processing will continue.", MsgBoxStyle.Critical + MsgBoxStyle.OkOnly)

                'Write to the error log, if enabled
                If My.MySettings.Default.errLog = True Then
                    errorText &= declist(j).ToString() & vbCrLf
                End If

                failed += 1
            End Try
            pbarMain.Increment(1)
        Next

        'If there were errors, write out the error log to a file; else delete unused variables
        If My.MySettings.Default.errLog = True And failed > 0 Then
            errorText &= "= = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = = =" & vbCrLf & vbCrLf & "END OF LOG"
            Computer.FileSystem.WriteAllText(desktop & "\" & "AES-ENCRYPTOR-ERROR_" & FileDate() & ".log", errorText, False)
        Else
            errorText = Nothing
            errorReport = Nothing
        End If
        Return
    End Sub

    Private Function GetTypeCount(ByVal type As Int32) As Int32
        Dim num As Int32
        Select Case type
            Case 0 'ENCRYPT
                For i As Int32 = 0 To (lvwLoad.Items.Count - 1) 'Iterate through all the process types, weeding out all but the ENCRYPTs
                    If lvwLoad.Items(i).SubItems(1).Text = "ENCRYPT" Then
                        num += 1
                    End If
                Next
            Case 1 'DECRYPT
                For i As Int32 = 0 To (lvwLoad.Items.Count - 1) 'Iterate through all the process types, weeding out all but the DECRYPTs
                    If lvwLoad.Items(i).SubItems(1).Text = "DECRYPT" Then
                        num += 1
                    End If
                Next
        End Select

        Return num
    End Function

    Private Sub btnOptions_Click(sender As System.Object, e As System.EventArgs) Handles btnOptions.Click
        dlgOptions.ShowDialog()
    End Sub

    Private Function FileDate() As String
        Return String.Format("{0:yyyy}-{0:MM}-{0:dd}_{0:HH}-{0:mm}-{0:ss}", DateTime.Now)
    End Function

    Private Sub GenLog()
        Dim filename As String = "AES-ENCRYPTOR-LOG_" & FileDate() & ".log"
        Dim logReport As New LogFile
        Computer.FileSystem.WriteAllText(desktop & "\" & filename, logReport.TransformText(), False)
        Return
    End Sub

    Private Sub btnProcFolder_Click(sender As System.Object, e As System.EventArgs) Handles btnProcFolder.Click
        If fbdSelect.ShowDialog() = Windows.Forms.DialogResult.OK Then
            Dim searchParam As String() = {"*.aesx"}
            Dim files
            Select Case MsgBox("Click 'Yes' for encryption and 'No' for decryption.", MsgBoxStyle.YesNoCancel)
                Case MsgBoxResult.Yes
                    If MySettings.Default.folderOption = True Then
                        files = My.Computer.FileSystem.GetFiles(fbdSelect.SelectedPath, FileIO.SearchOption.SearchAllSubDirectories)
                    Else
                        files = My.Computer.FileSystem.GetFiles(fbdSelect.SelectedPath, FileIO.SearchOption.SearchTopLevelOnly)
                    End If
                    Dim tmp(2) As String
                    Dim lvi As ListViewItem
                    For i As Int32 = 0 To (files.Count - 1) 'Add the ENCRYPT property to all the files within the specified folder
                        tmp(0) = files.Item(i.ToString)
                        tmp(1) = "ENCRYPT"
                        lvi = New ListViewItem(tmp)
                        lvwLoad.Items.Add(lvi)
                    Next

                Case MsgBoxResult.No
                    If MySettings.Default.folderOption = True Then
                        files = My.Computer.FileSystem.GetFiles(fbdSelect.SelectedPath, FileIO.SearchOption.SearchAllSubDirectories, searchParam)
                    Else
                        files = My.Computer.FileSystem.GetFiles(fbdSelect.SelectedPath, FileIO.SearchOption.SearchTopLevelOnly, searchParam)
                    End If
                    Dim tmp(2) As String
                    Dim lvi As ListViewItem
                    For i As Int32 = 0 To (files.Count - 1) 'Add the DECRYPT property to all the files within the specified folder
                        tmp(0) = files.Item(i.ToString)
                        tmp(1) = "DECRYPT"
                        lvi = New ListViewItem(tmp)
                        lvwLoad.Items.Add(lvi)
                    Next
                Case MsgBoxResult.Cancel
                    Exit Sub
            End Select
        End If

    End Sub

    Private Function LookForEncrypt() As Boolean
        Dim procList(lvwLoad.Items.Count) As String
        For i As Int32 = 0 To (lvwLoad.Items.Count - 1) 'Enumerate all the process types
            procList(i) = lvwLoad.Items(i).SubItems(1).Text
        Next

        If procList.Contains("ENCRYPT") Then
            Return True
        Else
            Return False
        End If
    End Function

    Private Function ObfucateFilename(ByVal oldFileName As String) As String
        Dim cryptor As New SimpleAes(MySettings.Default.password)
        Dim cryptedString = Convert.ToBase64String(cryptor.EncryptData(Encoding.UTF8.GetBytes(oldFileName)))
        cryptedString = cryptedString.Replace("/", "-")
        Return cryptedString
    End Function

    Private Function IsBase64(ByVal str As String) As Boolean
        Try
            Dim data As Byte() = Convert.FromBase64String(str.Replace("-", "/"))
            If str.Replace(" ", "").Length Mod 4 = 0 Then
                Return True
            Else : Return False
            End If
        Catch ex As Exception
            Return False
        End Try
    End Function

    Private Function DeobfuscateFilename(ByVal filename As String) As String
        Dim cryptor As New SimpleAes(MySettings.Default.password)

        Dim decryptedString = cryptor.DecryptData(Convert.FromBase64String(filename.Replace("-", "/")))
        Return Encoding.UTF8.GetChars(decryptedString)
    End Function
End Class